#!/usr/bin/python
'''
Created by Colin Gillespie <c.gillespie@ncl.ac.uk>

Last Changed on: $Date$ 
'''

__revision__ = '$Revision$'

import sys
try:
    import psyco
    psyco.full()
except ImportError:
    pass

#Non-standard modules
from mrange import cumRange
from VariableDict import Variables
from GetMomentEquation import getMomentEquation, getCumulantEquation
import Maple
from GetReactionInformation import GetStoichiometryDict
from GlobalVariables import globalVariables

class MapleInformation(object):
    '''Used as a storage container'''
    
    def __init__(self):
        self.pars = []
        self.species = []
        self.mom_equs = []

        
    def putFileName(self, file_name):
        self.file_name = file_name
        self.fp = open(self.file_name, 'w')
    
    def getFile(self):
        return self.fp

    def closeFile(self):
        self.fp.close()
        return    
    
    def putListOfSpecies(self, list_of_species):
        self.species = list_of_species
    
    def getListOfSpecies(self):
        return self.species
    
    def putStoicDict(self, gsd):
        self.gsd = gsd
     
    def getStoicDict(self):
        return self.gsd

    def putListOfPars(self, list_of_pars):
        self.pars = list_of_pars

    def getListOfPars(self):
        return self.pars
    
    def putListOfMomEqus(self, list_of_mom_equs):
        self.mom_equs = list_of_mom_equs
    
    def getListOfMomEqus(self):
        return self.mom_equs
    

def writeHeader(maple_information):
    '''Writes the maple file header.'''
    m = maple_information
    
    m.getFile().write(Maple.getHeader())
    expression = Maple.wrapExpression('restart;')
    m.getFile().write(expression)
    expression = Maple.wrapExpression('infolevel[dsolve]:=2;')
    m.getFile().write(expression)
    return True

def writeParameters(maple_information):
    m = maple_information
    
    for pa in m.getStoicDict().pa_dict.keys():
       # print '...Getting the parameter %s' % pa
        m.pars.append(pa + ':=' + str(m.getStoicDict().pa_dict[pa]['value']))
    

    #Get ready for outputing into maple   
    section_header = Maple.getSectionHeader('Constant Species and Parameters')
    m.getFile().write(section_header)
    maple_pars = ";".join(m.pars) + ";"    
    expression = Maple.wrapExpression(maple_pars)
    m.getFile().write(expression)
    section_footer = Maple.getSectionFooter()
    m.getFile().write(section_footer)
    return True

def writeDifferentialEquations(maple_information):
    m = maple_information
    
    section_header = Maple.getSectionHeader('ODEs')
    m.getFile().write(section_header)
    i = 0
    eq_list = []
    expression = ''
    for species, reaction in zip(m.species, m.mom_equs):
        expression += 'eq_%s:= diff(%s, t)= %s;\n' % (i, str(Variables().getMoment(species)), reaction)
        eq_list.append('eq_' + str(i))
        i += 1

    expression = Maple.wrapExpression(expression)
    m.getFile().write(expression)

    section_footer = Maple.getSectionFooter()
    m.getFile().write(section_footer)
    return tuple(eq_list)


def getEqualsTo(m, specie_tup):
    equals_to = ''
    for i, value in enumerate(specie_tup):
        if value != 0:
            sp = m.getStoicDict().m.getSpecies(i)
            if len(equals_to):
                equals_to += '*' + str(sp.getInitialAmount())
            else:
                equals_to = str(sp.getInitialAmount())
            if value > 1:
                equals_to += '**%s' % str(value)
    equals_to = '(0) = ' + equals_to
    return equals_to

def writeMapleSpeciesAtEvent(m, eq_tup, event_no, time, assignments, value):
    maple_species = []
    eq_list = list(eq_tup)
   
    for i, species in enumerate(m.species):
        t_species = tuple(species)
        if sum(species) > 1: #higher order moment
            if time == '0':
                equals_to = getEqualsTo(m, species)
                #equals_to = '(%s) = 0' % time
                maple_sp = str(Variables().getMoment(t_species))[:-3] + equals_to
            else:
                raise "Don't do events yet"
#                equals_to = '(%s) = h%s(%s)[%s]' % (time, str(int(event_no)-1), time, str(i+1))
#                for j,v in enumerate(species):
#                    if v > 0:
#                        sp = GDD.m.getSpecies(j)   
#                        if sp.getId() in assignments:
#                            equals_to = '(%s) = 0' % time
#                maple_sp = str(Variables(tuple(species), True).get())[:-3] + equals_to       
        else:
            index = list(species).index(1)
            sp = m.getStoicDict().m.getSpecies(index) 
            if sp.getId() in assignments:
                raise "Don't do events yet"
#                value_index = assignments.index(sp.getId())
#                equals_to = '(%s) = %s' % (str(time), str(value[value_index]))
#                maple_sp = str(Variables(tuple(species), True).get())[:-3] + equals_to
            else:
                if time == '0':
                    equals_to = '(%s) = %s' % (str(time), str(sp.getInitialAmount()))
                    maple_sp = str(Variables().getMoment(t_species))[:-3] + equals_to
                else:
                    raise "Don't do events yet"
#                    equals_to = '(%s) = h%s(%s)[%s]' %(time, str(int(event_no)-1), time, str(i+1))
#                    maple_sp = str(Variables(tuple(species), True).get())[:-3] + equals_to
        maple_species.append(str(Variables().getMoment(t_species)))
        eq_list.append(maple_sp)
    
    maple_initial_conditions = 'IC' + event_no + ':={' + ", ".join(eq_list) + '};'
    expression = Maple.wrapExpression(maple_initial_conditions)
    m.getFile().write(expression)

    maple_species = ",".join(maple_species)
    
    part =  'F' + event_no + ':=dsolve(IC' + event_no + ', {' + maple_species + '}, type=numeric, method=rosenbrock);'
    expression = Maple.wrapExpression(part)
    m.getFile().write(expression) 
    
    part = 'h' + event_no + ':= theta->eval([' + maple_species+ '],F' + event_no + '(theta));'
    expression = Maple.wrapExpression(part)
    m.getFile().write(expression)
     
    return True


def writeEvents(maple_information, eq_tup):
    m = maple_information
    writeMapleSpeciesAtEvent(m, eq_tup, '0', '0', [], [])
#    events, order = GetEvents(GDD.m).getEvents()
#    print events
#    no_of_events = len(events)
#    for i, event in enumerate(order):
#        time = events[event]['trigger'][2]
#        assignments = events[event]['assignments']
#        values = events[event]['values']
#        writeMapleSpeciesAtEvent(file_handler, list_of_species, GDD, eq_tup, str(i+1), time, assignments,values)

    

def createMapleFile(sbml_file, output_file=None):
    '''Creates the maple file with moment equations.
    
    Arguments:
        sbml_file - sbml file name
        output_file - outout file
    '''
    
    gsd = GetStoichiometryDict(sbml_file = sbml_file)
    m = MapleInformation()
    m.putFileName(output_file)
    m.putStoicDict(gsd)
    
    reaction_list = gsd.getReactionList()
    max_values = [globalVariables()['closure_order']] * len(gsd.empty_list)

    for values in cumRange(max_values):
        mom_eq = getMomentEquation(values, reaction_list)
        #print values
        if mom_eq != 0:
            m.species.append(values)
            #print '...Getting moment equation  dkappa_%s/dt' % str(values)
            m.mom_equs.append(str(mom_eq))
        else:
 #           print '0:=', values
            if sum(values) > 1: #higher order cumulant
                maple_sp = str(Variables().getMoment(tuple(values)))[:-3] + '(t):= 0'
            else:
                index = str(                index = list(values).index(1))
                raise 'Species %s is constant. Please add - constant ="true" - to your SBML file' % index
                index = list(values).index(1)
                sp = gsd.m.getSpecies(index)
                maple_sp = str(Variables().getMoment(tuple(values)))[:-3] + '(t):= ' + str(sp.getInitialAmount())
           
            m.pars.append(maple_sp)
            #print '***Species %s is zero' % str(Variables(values, True).get())
   

    writeHeader(m)
    writeParameters(m)
    eq_tup = writeDifferentialEquations(m)
    writeEvents(m, eq_tup) 
    
    m.getFile().write(Maple.getFooter())
    m.closeFile()
    print m.gsd.reaction_list[0].species_order
    return True



if __name__ == '__main__':
    print getDifferentialEquations(sbml_file='xml_files/imdeath.xml', output_file='/tmp/tmp.txt')

    #lambda_dict = GetLambdaDict(sbml_file='xml_files/sis.xml').GetLambdaDict()
#    print getDifferentialEquations(sbml_file='../xml_files/Hsp90model1.xml', moment_order = 2)
#    print getDifferentialEquations(sbml_file='xml_files/tmp.xml', moment_order = 2)
#    print getDifferentialEquations(sbml_file='xml_files/sis.xml', moment_order = 2)

    
    
    

    print 'EOF'
#h:=theta->eval(kappa_00000000000002000(t),F(theta));
 
#plot(h, 0..300);
